import os
import json
import datetime
import uuid
import numpy as np
import cv2
from flask import Blueprint, render_template, redirect, url_for, flash, request, current_app, jsonify, send_file
from flask_login import login_required, current_user
from models.database import db, Image, Prediction, AnalysisResult
from utils.analysis import calculate_area, generate_time_series, generate_thematic_map, export_to_excel
import matplotlib.pyplot as plt

# 定义一个处理datetime的JSON编码器
class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime.datetime):
            return obj.strftime('%Y-%m-%d %H:%M:%S')
        elif isinstance(obj, datetime.date):
            return obj.strftime('%Y-%m-%d')
        return super().default(obj)

analysis_bp = Blueprint('analysis', __name__, url_prefix='/analysis')

@analysis_bp.route('/area/<int:prediction_id>')
@login_required
def area_analysis(prediction_id):
    prediction = Prediction.query.get_or_404(prediction_id)
    image = Image.query.get_or_404(prediction.image_id)
    
    # 确保用户只能查看自己的分析结果
    if image.user_id != current_user.id:
        flash('您无权访问该资源')
        return redirect(url_for('data.my_data'))
    
    # 获取水域分割结果的路径
    result_path = os.path.join(current_app.config['RESULTS_FOLDER'], prediction.result_filename)
    
    # 计算水域面积
    area_data = calculate_area(result_path, image.latitude, image.longitude)
    
    # 保存分析结果
    result_data = json.dumps(area_data)
    
    # 检查是否已有此类型的分析结果
    existing_analysis = AnalysisResult.query.filter_by(
        prediction_id=prediction.id,
        result_type='area_calculation'
    ).first()
    
    if existing_analysis:
        # 更新现有结果
        existing_analysis.result_data = result_data
        db.session.commit()
        analysis_id = existing_analysis.id
    else:
        # 创建新的分析结果
        analysis_result = AnalysisResult(
            prediction_id=prediction.id,
            result_type='area_calculation',
            result_data=result_data
        )
        db.session.add(analysis_result)
        db.session.commit()
        analysis_id = analysis_result.id
    
    return render_template('analysis/area.html', 
                          prediction=prediction,
                          image=image,
                          area_data=area_data,
                          analysis_id=analysis_id)

@analysis_bp.route('/time-series')
@login_required
def time_series_analysis():
    # 获取用户的所有预测结果
    user_images = Image.query.filter_by(user_id=current_user.id).all()
    image_ids = [img.id for img in user_images]
    predictions = Prediction.query.filter(Prediction.image_id.in_(image_ids)).all()
    
    if not predictions:
        flash('没有足够的数据进行时间序列分析')
        return redirect(url_for('data.my_data'))
    
    # 按模型类型分组
    model_predictions = {
        'U²-Net (推荐)': [],
        'UNet': [],
        'DeepLabV3+': [],
        'MNDWI': []
    }
    
    # 添加备用键以处理模型名称可能存在的变体
    model_keys = {
        'U²-Net': 'U²-Net (推荐)',
        'u2net': 'U²-Net (推荐)',
        'UNet': 'UNet',
        'unet': 'UNet',
        'DeepLabV3+': 'DeepLabV3+',
        'deeplabv3': 'DeepLabV3+',
        'MNDWI': 'MNDWI',
        'mndwi': 'MNDWI'
    }
    
    for pred in predictions:
        # 确定模型键
        model_key = model_keys.get(pred.model_name, 'UNet')  # 默认为UNet
        
        # 确保键存在于model_predictions
        if model_key not in model_predictions:
            model_predictions[model_key] = []
            
        model_predictions[model_key].append({
            'id': pred.id,
            # 将datetime转换为字符串格式
            'date': pred.prediction_date.strftime('%Y-%m-%d %H:%M:%S'),
            'water_area': pred.water_area,
            'image_id': pred.image_id
        })
    
    # 生成时间序列图表
    chart_filename = generate_time_series(model_predictions)
    
    # 保存分析结果 - 使用自定义的JSON编码器
    result_data = json.dumps(model_predictions, cls=DateTimeEncoder)
    analysis_result = AnalysisResult(
        prediction_id=predictions[0].id,  # 使用第一个预测ID作为参考
        result_type='time_series',
        result_data=result_data,
        chart_filename=chart_filename
    )
    
    db.session.add(analysis_result)
    db.session.commit()
    
    return render_template('analysis/time_series.html',
                          chart_url=url_for('static', filename=f'results/{chart_filename}'),
                          model_predictions=model_predictions,
                          analysis_id=analysis_result.id)

@analysis_bp.route('/thematic-map/<int:prediction_id>')
@login_required
def thematic_map(prediction_id):
    prediction = Prediction.query.get_or_404(prediction_id)
    image = Image.query.get_or_404(prediction.image_id)
    
    # 确保用户只能查看自己的分析结果
    if image.user_id != current_user.id:
        flash('您无权访问该资源')
        return redirect(url_for('data.my_data'))
    
    # 获取水域分割结果的路径
    result_path = os.path.join(current_app.config['RESULTS_FOLDER'], prediction.result_filename)
    
    # 生成专题图
    map_filename = generate_thematic_map(result_path, image.latitude, image.longitude, image.city)
    
    # 保存分析结果
    existing_analysis = AnalysisResult.query.filter_by(
        prediction_id=prediction.id,
        result_type='thematic_map'
    ).first()
    
    if existing_analysis:
        # 更新现有结果
        existing_analysis.chart_filename = map_filename
        db.session.commit()
        analysis_id = existing_analysis.id
    else:
        # 创建新的分析结果
        analysis_result = AnalysisResult(
            prediction_id=prediction.id,
            result_type='thematic_map',
            chart_filename=map_filename
        )
        db.session.add(analysis_result)
        db.session.commit()
        analysis_id = analysis_result.id
    
    return render_template('analysis/thematic_map.html',
                          prediction=prediction,
                          image=image,
                          map_url=url_for('static', filename=f'results/{map_filename}'),
                          analysis_id=analysis_id)

@analysis_bp.route('/export-excel/<int:prediction_id>')
@login_required
def export_excel(prediction_id):
    prediction = Prediction.query.get_or_404(prediction_id)
    image = Image.query.get_or_404(prediction.image_id)
    
    # 确保用户只能导出自己的数据
    if image.user_id != current_user.id:
        flash('您无权访问该资源')
        return redirect(url_for('data.my_data'))
    
    # 获取水域分割结果的路径
    result_path = os.path.join(current_app.config['RESULTS_FOLDER'], prediction.result_filename)
    
    # 生成Excel文件
    excel_path = export_to_excel(result_path, image, prediction)
    
    # 返回文件供下载
    return send_file(excel_path, as_attachment=True, download_name=f'water_analysis_{prediction.id}.xlsx') 